"""
记忆管理器模块
统一管理传统记忆和增强记忆系统，提供智能降级机制

作者: 架构重构优化
日期: 2025-01-05
"""

import asyncio
import logging
from typing import Dict, Any, List, Optional
from langchain_core.messages import HumanMessage, AIMessage, BaseMessage
from .Memory import MemoryClass
from .config_manager import ConfigManager
from .async_manager import get_async_manager, run_async

logger = logging.getLogger(__name__)

class MemoryManager:
    """
    统一记忆管理器
    支持传统记忆和增强记忆的智能切换
    """
    
    def __init__(self, config: ConfigManager):
        """
        初始化记忆管理器
        
        Args:
            config: 配置管理器实例
        """
        self.config = config
        self.enhanced_mode = config.memory.enhanced_memory_enabled
        self.traditional_memory = None
        self.enhanced_memory = None
        self._enhanced_available = False
        
        # 初始化传统记忆系统
        self._init_traditional_memory()
        
        # 尝试初始化增强记忆系统
        if self.enhanced_mode:
            self._init_enhanced_memory()
        
        logger.info(f"记忆管理器初始化完成，增强模式: {self._enhanced_available}")

    def _init_traditional_memory(self):
        """初始化传统记忆系统"""
        try:
            self.traditional_memory = MemoryClass(
                memoryKey=self.config.memory.memory_key,
                model=self.config.model.base_model
            )
            logger.info("传统记忆系统初始化成功")
        except Exception as e:
            logger.error(f"传统记忆系统初始化失败: {e}")
            raise

    def _init_enhanced_memory(self):
        """初始化增强记忆系统"""
        try:
            # 尝试导入增强记忆模块
            from enhanced_memory.enhanced_memory_class import EnhancedMemoryClass
            
            self.enhanced_memory = EnhancedMemoryClass()
            self._enhanced_available = True
            logger.info("增强记忆系统初始化成功")
            
        except ImportError as e:
            logger.warning(f"增强记忆模块不可用: {e}")
            self._enhanced_available = False
        except Exception as e:
            logger.error(f"增强记忆系统初始化失败: {e}")
            self._enhanced_available = False

    def get_memory(self, session_id: str) -> Any:
        """
        获取会话记忆
        
        Args:
            session_id: 会话ID
            
        Returns:
            记忆历史对象
        """
        try:
            # 优先使用增强记忆
            if self._enhanced_available and self.enhanced_memory:
                try:
                    # 增强记忆可能需要异步处理
                    return self._get_enhanced_memory_sync(session_id)
                except Exception as e:
                    logger.warning(f"增强记忆获取失败，降级到传统记忆: {e}")
            
            # 使用传统记忆
            return self.traditional_memory.get_memory(session_id)
            
        except Exception as e:
            logger.error(f"记忆获取失败: {e}")
            # 返回空的记忆对象
            return self._create_empty_memory()

    def _get_enhanced_memory_sync(self, session_id: str) -> Any:
        """同步方式获取增强记忆（使用异步管理器简化）"""
        try:
            # 使用异步管理器简化异步操作
            async def get_enhanced():
                if hasattr(self.enhanced_memory, 'get_conversation_history'):
                    return await self.enhanced_memory.get_conversation_history(session_id)
                else:
                    # 如果增强记忆没有对应方法，降级到传统记忆
                    return self.traditional_memory.get_memory(session_id)

            result = run_async(get_enhanced(), timeout=10.0)
            return result if result is not None else self.traditional_memory.get_memory(session_id)

        except Exception as e:
            logger.warning(f"增强记忆同步获取失败: {e}")
            return self.traditional_memory.get_memory(session_id)

    def save_memory(self, session_id: str, user_input: str, ai_output: str):
        """
        保存对话到记忆
        
        Args:
            session_id: 会话ID
            user_input: 用户输入
            ai_output: AI输出
        """
        try:
            # 保存到传统记忆（始终保存，作为备份）
            self._save_to_traditional_memory(session_id, user_input, ai_output)
            
            # 保存到增强记忆
            if self._enhanced_available and self.enhanced_memory:
                try:
                    self._save_to_enhanced_memory(session_id, user_input, ai_output)
                except Exception as e:
                    logger.warning(f"增强记忆保存失败: {e}")
            
        except Exception as e:
            logger.error(f"记忆保存失败: {e}")

    def _save_to_traditional_memory(self, session_id: str, user_input: str, ai_output: str):
        """保存到传统记忆"""
        try:
            memory_history = self.traditional_memory.get_memory(session_id)
            
            # 添加用户消息
            user_message = HumanMessage(content=user_input)
            memory_history.add_message(user_message)
            
            # 添加AI回复
            ai_message = AIMessage(content=ai_output)
            memory_history.add_message(ai_message)
            
            # 保存记忆（如果有save_memory方法）
            if hasattr(self.traditional_memory, 'save_memory'):
                self.traditional_memory.save_memory(session_id, memory_history)
            logger.debug(f"传统记忆保存成功: {session_id}")
            
        except Exception as e:
            logger.error(f"传统记忆保存失败: {e}")

    def _save_to_enhanced_memory(self, session_id: str, user_input: str, ai_output: str):
        """保存到增强记忆，使用同步Redis存储避免事件循环问题"""
        try:
            # 直接使用传统记忆存储，避免异步Redis操作
            logger.debug(f"使用传统记忆存储避免Redis事件循环问题: {session_id}")

            # 使用传统记忆系统存储
            if hasattr(self, 'traditional_memory') and self.traditional_memory:
                self.traditional_memory.save_memory(session_id, user_input, ai_output)
                logger.debug(f"传统记忆保存成功: {session_id}")
            else:
                logger.warning(f"传统记忆系统不可用: {session_id}")

        except Exception as e:
            logger.warning(f"记忆保存失败: {e}")

    def clear_memory(self, session_id: str) -> bool:
        """
        清除会话记忆
        
        Args:
            session_id: 会话ID
            
        Returns:
            是否清除成功
        """
        success = True
        
        try:
            # 清除传统记忆
            memory_history = self.traditional_memory.get_memory(session_id)
            if hasattr(memory_history, 'clear'):
                memory_history.clear()
            logger.debug(f"传统记忆清除成功: {session_id}")
            
        except Exception as e:
            logger.error(f"传统记忆清除失败: {e}")
            success = False
        
        # 清除增强记忆
        if self._enhanced_available and self.enhanced_memory:
            try:
                self._clear_enhanced_memory(session_id)
                logger.debug(f"增强记忆清除成功: {session_id}")
            except Exception as e:
                logger.warning(f"增强记忆清除失败: {e}")
                # 增强记忆清除失败不影响整体结果
        
        return success

    def _clear_enhanced_memory(self, session_id: str):
        """清除增强记忆（使用异步管理器简化）"""
        try:
            if hasattr(self.enhanced_memory, 'clear_conversation'):
                async def clear_enhanced():
                    return await self.enhanced_memory.clear_conversation(session_id)

                # 使用异步管理器执行清除操作
                run_async(clear_enhanced(), timeout=5.0)
        except Exception as e:
            logger.warning(f"增强记忆清除操作失败: {e}")

    def get_memory_context(self, session_id: str, max_messages: int = 6) -> str:
        """
        获取记忆上下文字符串
        
        Args:
            session_id: 会话ID
            max_messages: 最大消息数量
            
        Returns:
            格式化的记忆上下文
        """
        try:
            memory_history = self.get_memory(session_id)
            
            if memory_history and hasattr(memory_history, 'messages') and memory_history.messages:
                # 获取最近的消息
                recent_messages = memory_history.messages[-max_messages:]
                
                context_parts = []
                for msg in recent_messages:
                    if hasattr(msg, 'content'):
                        content = msg.content
                        # 限制每条消息的长度
                        if len(content) > 100:
                            content = content[:100] + "..."
                        
                        if isinstance(msg, HumanMessage):
                            context_parts.append(f"用户: {content}")
                        elif isinstance(msg, AIMessage):
                            context_parts.append(f"助手: {content}")
                
                if context_parts:
                    # 保留最近的对话
                    context = " | ".join(context_parts[-4:])
                    return f"之前的对话: {context}"
            
            return ""
            
        except Exception as e:
            logger.warning(f"获取记忆上下文失败: {e}")
            return ""

    def _create_empty_memory(self):
        """创建空的记忆对象"""
        class EmptyMemory:
            def __init__(self):
                self.messages = []
            
            def add_message(self, message):
                pass
            
            def clear(self):
                pass
        
        return EmptyMemory()

    def get_stats(self) -> Dict[str, Any]:
        """获取记忆管理器统计信息"""
        return {
            "enhanced_mode": self.enhanced_mode,
            "enhanced_available": self._enhanced_available,
            "traditional_memory_available": self.traditional_memory is not None,
            "memory_key": self.config.memory.memory_key,
            "max_token_limit": self.config.memory.max_token_limit
        }
